package com.gandong8.trace.tool.util;

import io.netty.util.concurrent.DefaultThreadFactory;
import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.*;

//Executor：是Java线程池的超级接口
//ExecutorService：它是线程池定义的一个接口，继承Executor。有两个实现类，分别为ThreadPoolExecutor,ScheduledThreadPoolExecutor
//ThreadPoolTaskExecutor是spring core包中的，而ThreadPoolExecutor是JDK中的JUC。ThreadPoolTaskExecutor是对ThreadPoolExecutor进行了封装处理
//Executors: 为线程迟工具类，相当于一个工厂类，用来创建合适的线程池，返回ExecutorService类型的线程池
public class ThreadUtil {

    public static Executor initExecutor(int corePoolSize, int maxPoolSize, int keepAliveSeconds,
                                         int queueCapacity, ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor() {
            @Override
            public <T> Future<T> submit(Callable<T> task) {
                return super.submit(ThreadUtil.wrap(task, MDC.getCopyOfContextMap()));
            }
            @Override
            public void execute(Runnable task) {
                super.execute(ThreadUtil.wrap(task, MDC.getCopyOfContextMap()));
            }
        };
        executor.setCorePoolSize(corePoolSize);
        executor.setMaxPoolSize(maxPoolSize);
        executor.setKeepAliveSeconds(keepAliveSeconds);
        executor.setThreadFactory(threadFactory);
        executor.setQueueCapacity(queueCapacity);
        executor.setRejectedExecutionHandler(rejectedExecutionHandler);
        executor.initialize();
        return executor;
    }

    public static Executor fixedExecutor(int num) {
        return fixedExecutor(num, null);
    }

    public static Executor fixedExecutor(int num, String threadName) {
        ThreadFactory threadFactory = threadName == null ? Executors.defaultThreadFactory():new DefaultThreadFactory(threadName);
        return initExecutor(num, num, 0, 10, threadFactory, new ThreadPoolExecutor.AbortPolicy());
    }

    public static Executor singleExecutor(String threadName) {
        return fixedExecutor(1, threadName);
    }

    public static Executor singleExecutor() {
        return singleExecutor(null);
    }

    private static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            TraceUtil.getTraceId();
            try {
                return callable.call();
            } finally {
                MDC.clear();
            }
        };
    }

    private static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            TraceUtil.getTraceId();
            try {
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }



}
